package pl;

import javacard.framework.ISO7816;
import javacard.framework.ISOException;
import javacard.framework.JCSystem;
import javacard.framework.Util;
import javacard.security.MessageDigest;
import javacard.security.RandomData;

public class EMSA_PSS {

	private static final byte _0X_BC = (byte) 0xBC;

	// https://www.das-labor.org/trac/browser/microcontroller-2/arm-crypto-lib/testvectors/rsa-pkcs-1v2-1-vec/pss-vect.txt
	// https://github.com/bdauvergne/python-pkcs1/blob/master/tests/data/pss-int.txt
	// http://grepcode.com/file/repo1.maven.org/maven2/org.bouncycastle/bcprov-jdk15on/1.47/org/bouncycastle/crypto/signers/PSSSigner.java

	/*
	 * byte[] message = new byte[] { (byte) 0x85, (byte) 0x9E, (byte) 0xEF,
	 * (byte) 0x2F, (byte) 0xD7, (byte) 0x8A, (byte) 0xCA, (byte) 0x00, (byte)
	 * 0x30, (byte) 0x8B, (byte) 0xDC, (byte) 0x47, (byte) 0x11, (byte) 0x93,
	 * (byte) 0xBF, (byte) 0x55, (byte) 0xBF, (byte) 0x9D, (byte) 0x78, (byte)
	 * 0xDB, (byte) 0x8F, (byte) 0x8A, (byte) 0x67, (byte) 0x2B, (byte) 0x48,
	 * (byte) 0x46, (byte) 0x34, (byte) 0xF3, (byte) 0xC9, (byte) 0xC2, (byte)
	 * 0x6E, (byte) 0x64, (byte) 0x78, (byte) 0xAE, (byte) 0x10, (byte) 0x26,
	 * (byte) 0x0F, (byte) 0xE0, (byte) 0xDD, (byte) 0x8C, (byte) 0x08, (byte)
	 * 0x2E, (byte) 0x53, (byte) 0xA5, (byte) 0x29, (byte) 0x3A, (byte) 0xF2,
	 * (byte) 0x17, (byte) 0x3C, (byte) 0xD5, (byte) 0x0C, (byte) 0x6D, (byte)
	 * 0x5D, (byte) 0x35, (byte) 0x4F, (byte) 0xEB, (byte) 0xF7, (byte) 0x8B,
	 * (byte) 0x26, (byte) 0x02, (byte) 0x1C, (byte) 0x25, (byte) 0xC0, (byte)
	 * 0x27, (byte) 0x12, (byte) 0xE7, (byte) 0x8C, (byte) 0xD4, (byte) 0x69,
	 * (byte) 0x4C, (byte) 0x9F, (byte) 0x46, (byte) 0x97, (byte) 0x77, (byte)
	 * 0xE4, (byte) 0x51, (byte) 0xE7, (byte) 0xF8, (byte) 0xE9, (byte) 0xE0,
	 * (byte) 0x4C, (byte) 0xD3, (byte) 0x73, (byte) 0x9C, (byte) 0x6B, (byte)
	 * 0xBF, (byte) 0xED, (byte) 0xAE, (byte) 0x48, (byte) 0x7F, (byte) 0xB5,
	 * (byte) 0x56, (byte) 0x44, (byte) 0xE9, (byte) 0xCA, (byte) 0x74, (byte)
	 * 0xFF, (byte) 0x77, (byte) 0xA5, (byte) 0x3C, (byte) 0xB7, (byte) 0x29,
	 * (byte) 0x80, (byte) 0x2F, (byte) 0x6E, (byte) 0xD4, (byte) 0xA5, (byte)
	 * 0xFF, (byte) 0xA8, (byte) 0xBA, (byte) 0x15, (byte) 0x98, (byte) 0x90,
	 * (byte) 0xFC };
	 */
	byte[] saltV = new byte[] { (byte) 0xE3, (byte) 0xB5, (byte) 0xD5,
			(byte) 0xD0, (byte) 0x02, (byte) 0xC1, _0X_BC, (byte) 0xE5,
			(byte) 0x0C, (byte) 0x2B, (byte) 0x65, (byte) 0xEF, (byte) 0x88,
			(byte) 0xA1, (byte) 0x88, (byte) 0xD8, (byte) 0x3B, (byte) 0xCE,
			(byte) 0x7E, (byte) 0x61 };

	byte[] hashSHA1 = new byte[] { (byte) 0x37, (byte) 0xB6, (byte) 0x6A,
			(byte) 0xE0, (byte) 0x44, (byte) 0x58, (byte) 0x43, (byte) 0x35,
			(byte) 0x3D, (byte) 0x47, (byte) 0xEC, (byte) 0xB0, (byte) 0xB4,
			(byte) 0xFD, (byte) 0x14, (byte) 0xC1, (byte) 0x10, (byte) 0xE6,
			(byte) 0x2D, (byte) 0x6A };

	private static final byte[] ZERO_ARRAY = new byte[8];

	protected MessageDigest md;
	protected MessageDigest mgfDig;
	private short hLen;
	private RandomData rd;

	private byte[] hashResult;

	byte[] salt;
	byte[] result;
	byte[] dbMask;
	byte[] DB;

	byte[] C;
	byte[] hash;

	private short emLen;
	private short emBits;
	private short sLen;

	private short counter;
	private short hLen2;

	public EMSA_PSS(short saltLen, short modLen) {

		// default - sha1
		md = MessageDigest.getInstance(MessageDigest.ALG_SHA, false);

		hLen = (short) md.getLength();
		hashResult = JCSystem.makeTransientByteArray(hLen,
				JCSystem.CLEAR_ON_DESELECT);

		hash = JCSystem
				.makeTransientByteArray(hLen, JCSystem.CLEAR_ON_DESELECT);

		rd = RandomData.getInstance(RandomData.ALG_PSEUDO_RANDOM);

		sLen = saltLen;

		salt = JCSystem
				.makeTransientByteArray(sLen, JCSystem.CLEAR_ON_DESELECT);

		C = JCSystem.makeTransientByteArray((short) 4,
				JCSystem.CLEAR_ON_DESELECT);

		emBits = (short) (modLen - 1);

		// equals to module length
		emLen = (short) ((short) (emBits + 7) / 8);

		result = JCSystem.makeTransientByteArray(emLen,
				JCSystem.CLEAR_ON_DESELECT);

		DB = JCSystem.makeTransientByteArray((short) (emLen - hLen - 1),
				JCSystem.CLEAR_ON_DESELECT);

		// default for mgf1
		mgfDig = MessageDigest.getInstance(MessageDigest.ALG_SHA, false);

		dbMask = new byte[(short) (emLen - hLen - 1)];
		
		//emsaPssEncode(hashSHA1);
		
		//this.emsaPssVerify(this.result);

	}

	public void emsaPssEncode(byte[] hashedMessage) {

		//rd.generateData(salt, (short) 0, sLen);
		
		Util.arrayCopy(saltV, (short) 0, salt, (short) 0, (short) salt.length);

		if (emBits < (short) (8 * hLen + 8 * sLen + 9)) {
			/* TODO: Exception */
			ISOException.throwIt(ISO7816.SW_WRONG_LENGTH);
		}

		// hash (00 00 00 00 00 00 00 00 || hash || salt)
		md.reset();
		md.update(ZERO_ARRAY, (short) 0, (short) ZERO_ARRAY.length);
		md.update(hashedMessage, (short) 0, (short) hashedMessage.length);

		// add salt
		md.doFinal(salt, (short) 0, (short) salt.length, hashResult, (short) 0);

		// DB = PS || 01 || salt.
		// DB = null;
		// DB = new byte[(short) (emLen - hLen - 1)];

		ClearArray(DB);

		DB[(short) (emLen - sLen - hLen - 2)] = 0x01;

		Util.arrayCopy(salt, (short) 0, DB, (short) (emLen - sLen - hLen - 1),
				sLen);

		// mgf1
		// dbMask = null;
		// dbMask = new byte[(short) (emLen - hLen - 1)];
		// dbMask = mgf1(hashResult, (short) (emLen - hLen - 1));

		mgf1(hashResult, (short)0, (short)hashResult.length);

		for (short i = 0; i < DB.length; i++)
			DB[i] = (byte) (DB[i] ^ dbMask[i]);

		DB[0] &= (0xFF >>> (8 * emLen - emBits));

		Util.arrayCopy(DB, (short) 0, result, (short) 0,
				(short) (emLen - hLen - 1));
		Util.arrayCopy(hashResult, (short) 0, result,
				(short) (emLen - hLen - 1), hLen);

		result[(short) (emLen - 1)] = _0X_BC;
	}

	public boolean emsaPssVerify(byte[] decodedHashMessage)
	{
		Util.arrayCopy(decodedHashMessage, (short)0, result, (short)0, (short)decodedHashMessage.length);
		
		if ((byte)(result[(short)(result.length - 1)]) !=(_0X_BC))
		{
			ISOException.throwIt(ISO7816.SW_COMMAND_CHAINING_NOT_SUPPORTED);
		}
		
		mgf1(result, (short) (result.length - hLen - 1), hLen);
		
		// byte[] dbMask = maskGeneratorFunction1(result, result.length - hLen - 1, hLen,  result.length - hLen - 1);
		
		for (short i = 0; i != (short)dbMask.length; i++)
		{
			result[i] ^= dbMask[i];
		}
		
		result[0] &= (0xff >> ((result.length * 8) - emBits));
		
		for (short i = 0; i != (short)(result.length - hLen - sLen - 2); i++)
		{
			if (result[i] != 0)
			{
				//clearBlock(block);
				return false;
			}
		}
		
		if (result[(short)(result.length - hLen - sLen - 2)] != 0x01)
		{
			//clearBlock(block);
			return false;
		}

		//Util.arrayCopy(result, (short)(result.length - sLen - hLen - 1), hashResult, arg3, sLen);
		//System.arraycopy(result, result.length - sLen - hLen - 1, hash, mDash.length - sLen, sLen);

		return true;
		
	}
	
	/*
	 * Works OK private byte[] mgf1(byte[] seed, short length) {
	 * 
	 * mask = null; mask = new byte[length]; //
	 * JCSystem.makeTransientByteArray(length, counter = 0; hLen2 =
	 * mgfDig.getLength();
	 * 
	 * byte[] hashBuf = new byte[hLen2]; mgfDig.reset();
	 * 
	 * while (counter < (short) (length / hLen2)) { I2OSP(counter, C);
	 * mgfDig.update(seed, (short) 0, (short) seed.length); mgfDig.doFinal(C,
	 * (short) 0, (short) C.length, hash, (short) 0); Util.arrayCopy(hash,
	 * (short) 0, mask, (short) (counter * hLen2), hLen2); counter++; }
	 * 
	 * if ((short) (counter * hLen2) < length) { I2OSP(counter, C);
	 * mgfDig.update(seed, (short) 0, (short) seed.length); mgfDig.doFinal(C,
	 * (short) 0, (short) C.length, hash, (short) 0); Util.arrayCopy(hash,
	 * (short) 0, mask, (short) (counter * hLen2), (short) (mask.length -
	 * (short) (counter * hLen2))); }
	 * 
	 * return mask; }
	 */

	private void mgf1(byte[] seed, short offset, short seedLen) {

		ClearArray(dbMask);
		ClearArray(hash);

		counter = 0;
		hLen2 = mgfDig.getLength();

		mgfDig.reset();

		while (counter < (short) ((short) dbMask.length / hLen2)) {
			I2OSP(counter, C);
			mgfDig.update(seed, offset, seedLen);
			mgfDig.doFinal(C, (short) 0, (short) C.length, hash, (short) 0);
			Util.arrayCopy(hash, (short) 0, dbMask, (short) (counter * hLen2),
					hLen2);
			counter++;
		}

		if ((short) (counter * hLen2) < (short) dbMask.length) {
			I2OSP(counter, C);
			mgfDig.update(seed, offset, seedLen);
			mgfDig.doFinal(C, (short) 0, (short) C.length, hash, (short) 0);
			Util.arrayCopy(hash, (short) 0, dbMask, (short) (counter * hLen2),
					(short) (dbMask.length - (short) (counter * hLen2)));
		}
	}

	private void I2OSP(short i, byte[] sp) {
		sp[0] = (byte) (i >>> 24);
		sp[1] = (byte) (i >>> 16);
		sp[2] = (byte) (i >>> 8);
		sp[3] = (byte) (i >>> 0);
	}

	public static void ClearArray(byte[] array) {
		for (short i = 0; i < (short) array.length; i++) {
			array[i] = (byte) 0x00;
		}
	}
}
